#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IK数据重放脚本
用于重放并验证IK数据中的sol_q是否正确
"""

import os
import time
import pickle
import numpy as np
import argparse
from tqdm import tqdm
import sys
import os

# 添加项目根目录到Python路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
sys.path.insert(0, project_root)

# 添加teleop目录到Python路径
teleop_dir = os.path.join(project_root, 'teleop')
sys.path.insert(0, teleop_dir)

from teleop.utils.mat_tool import print_pose_info
from teleop.robot_control.fangnuo_car import FangnuoCar

def load_ik_data(file_path):
    """加载IK数据文件"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"IK数据文件不存在: {file_path}")
    
    with open(file_path, 'rb') as f:
        data = pickle.load(f)
    
    print(f"成功加载IK数据，共 {len(data)} 帧")
    return data

def replay_ik_data(ik_data, use_real_arms=False, speed=1.0):
    """
    重放IK数据
    
    参数:
        ik_data: IK数据列表
        use_real_arms: 是否使用真实机械臂
        speed: 重放速度倍数
    """
    print(f"初始化机械臂控制器 (使用{'真实' if use_real_arms else '模拟'}机械臂)...")
    car = FangnuoCar(use_real_arms=use_real_arms)
    
    # 使能机械臂
    if use_real_arms:
        print("使能机械臂...")
        car.enable_dual_arm()
    
    # 计算时间间隔
    if len(ik_data) > 1:
        # 假设数据是按固定时间间隔记录的
        dt = (ik_data[1]['timestamp'] - ik_data[0]['timestamp']) / speed
    else:
        dt = 0.1  # 默认100ms
    
    print(f"开始重放IK数据，共 {len(ik_data)} 帧，时间间隔: {dt*1000:.1f}ms")
    
    try:
        # 遍历每一帧数据
        for i, frame in enumerate(tqdm(ik_data, desc="重放进度")):
            if 'sol_q' not in frame or frame['sol_q'] is None:
                print(f"警告: 第 {i} 帧没有有效的sol_q数据，跳过")
                continue
                
            # 获取目标关节角度
            q_target = frame['sol_q']
            # q_target[4] = -q_target[4]
            # q_target[10] = -q_target[10]
            print(f"Left wrist pose: {frame['left_wrist']}")
            
            # 控制机械臂
            car.ctrl_dual_arm(q_target)
            # 获取当前机械臂状态
            status = car.status()
            
            print(f"Left arm position: {status['left_arm']['position']}")
            print(f"Right arm position: {status['right_arm']['position']}")
          
            
            # 等待到下一帧
            time.sleep(dt)
            
    except KeyboardInterrupt:
        print("\n用户中断重放")
    except Exception as e:
        print(f"重放过程中发生错误: {str(e)}")
    finally:
        print("重放结束")
        # 可以在这里添加机械臂归位或安全停止的代码

def main():
    parser = argparse.ArgumentParser(description='重放IK数据')
    parser.add_argument('--ik-file', type=str, default='data/ik_data/ik_data_20250605_161826.pkl', 
                       help='IK数据文件路径 (默认: data/ik_data)')
    parser.add_argument('--real', action='store_true', help='使用真实机械臂')
    parser.add_argument('--speed', type=float, default=1.0, help='重放速度倍数 (默认: 1.0)')
    
    args = parser.parse_args()
    
    try:
        # 加载IK数据
        ik_data = load_ik_data(args.ik_file)
        
        # 重放数据
        replay_ik_data(ik_data, use_real_arms=args.real, speed=args.speed)
        
    except Exception as e:
        print(f"错误: {str(e)}")
        return 1
    
    return 0

if __name__ == "__main__":
    exit(main())
